package cn.flying.cloud.base.dubbo.loadbalancer;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.dubbo.common.URL;
import org.apache.dubbo.common.extension.ExtensionLoader;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.Invoker;
import org.apache.dubbo.rpc.cluster.LoadBalance;
import org.apache.dubbo.rpc.cluster.loadbalance.AbstractLoadBalance;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;

import static org.apache.dubbo.common.constants.CommonConstants.DEFAULT_KEY_PREFIX;

import cn.flying.cloud.base.common.utils.Ft;

/**
 * 自定义负载均衡策略
 *
 * @author: admin
 * @date: 2025年03月25日 23:59
 * @version: 1.0
 */
public class CustomLoadBalance extends AbstractLoadBalance {
    private final Logger logger = LoggerFactory.getLogger(this.getClass());
    private static final String DEFAULT_STRATEGY = "roundrobin"; // 默认回退策略

    @Override
    protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation invocation) {
        // 从请求上下文获取
        String requestId = Ft.Local.RequestId();
        String envId = Ft.Local.EnvId();
        if (Ft.String.isBlank(requestId) || Ft.String.isBlank(envId)) {
            // 如果未携带 request-id，则直接使用默认策略
            return selectDefault(invokers, url, invocation);
        }

        Map<String, Invoker<T>> invokerMap = new HashMap<>();
        for (Invoker<T> invoker : invokers) {
            String key = invoker.getUrl().getHost() + ":" + invoker.getUrl().getPort();
            invokerMap.put(key, invoker);
        }

        // 获取服务实例列表
        String serviceName = url.getParameter("remote.application");
        DiscoveryClient discoveryClient = Ft.Spring.getBean(DiscoveryClient.class);
        List<ServiceInstance> instances = discoveryClient.getInstances(serviceName);
        // 根据 request-id 匹配服务实例
        for (ServiceInstance instance : instances) {
            Map<String, String> metadata = instance.getMetadata();
            String instanceRequestId = metadata.get("request-id");
            String instanceEnvId = metadata.get("env-id");
            if (requestId.equals(instanceRequestId) && envId.equals(instanceEnvId)) {
                // 找到对应instance的 Invoker
                String urls = metadata.get("dubbo.metadata-service.urls");
                Map<String, String> map = toMap(urls);
                String instancePort = map.get("bind.port");
                String instanceHost = map.get("bind.ip");

                Invoker<T> invoker = invokerMap.get(instanceHost + ":" + instancePort);
                if (invoker != null) {
                    logger.info("调用【{}】的实例，主机：{}端口：{}", instanceRequestId, instanceHost, instancePort);
                    return invoker;
                }
            }
        }

        // 如果没有匹配的实例，使用默认的负载均衡策略
        return selectDefault(invokers, url, invocation);
    }

    // 调用默认负载均衡策略（如 roundrobin）
    private <T> Invoker<T> selectDefault(List<Invoker<T>> invokers, URL url, Invocation invocation) {
        // 获取默认策略的 LoadBalance 实现
        LoadBalance defaultStrategy = ExtensionLoader.getExtensionLoader(LoadBalance.class).getExtension(DEFAULT_STRATEGY);
        return defaultStrategy.select(invokers, url, invocation);
    }

    private Map<String, String> toMap(String url) {
        url = url.substring(3, url.length() - 2);
        Map<String, String> parameters = null;
        int i = url.indexOf('?'); // separator between body and parameters
        if (i >= 0) {
            String[] parts = url.substring(i + 1).split("&");
            parameters = new HashMap<>();
            for (String part : parts) {
                part = part.trim();
                if (!part.isEmpty()) {
                    int j = part.indexOf('=');
                    if (j >= 0) {
                        String key = part.substring(0, j);
                        String value = part.substring(j + 1);
                        parameters.put(key, value);
                        // compatible with lower versions registering "default." keys
                        if (key.startsWith(DEFAULT_KEY_PREFIX)) {
                            parameters.putIfAbsent(key.substring(DEFAULT_KEY_PREFIX.length()), value);
                        }
                    } else {
                        parameters.put(part, part);
                    }
                }
            }
        }
        return parameters;
    }
}    